/**
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.hyperledger.stresstest.helper

import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler}
import akka.stream.{Attributes, BidiShape, Inlet, Outlet}

class PassivePingPongStage[M](f: PartialFunction[M, M]) extends GraphStage[BidiShape[M, M, M, M]] {
  val protoIn = Inlet[M]("protoIn")
  val protoOut = Outlet[M]("protoOut")
  val appIn = Inlet[M]("appIn")
  val appOut = Outlet[M]("appOut")

  override def createLogic(inheritedAttributes: Attributes) = new GraphStageLogic(shape) {
    override def preStart(): Unit = {
      tryPull(appIn)
      tryPull(protoIn)
    }

    passAlong(appIn, protoOut)
    setHandler(protoOut, eagerTerminateOutput)
    setHandler(appOut, eagerTerminateOutput)
    setHandler(protoIn, new InHandler with (() => Unit){
      def apply: Unit = tryPull(protoIn)
      override def onPush(): Unit = {
        val elem = grab(protoIn)
        if (f.isDefinedAt(elem))
          emit(protoOut, f(elem), this)
        else
          emit(appOut, elem, this)
      }
    })

  }
  override val shape = BidiShape(protoIn, appOut, appIn, protoOut)
}
